package com.margo.project.web.mvcconfig.interceptor;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 复制 com.netflix.discovery.util.RateLimiter
 * 	限流
 *  https://segmentfault.com/a/1190000013127775
 *  burst 突发
 * @author fangxingbang
 *
 */
public class MargoEurekaRateLimiter {
	private final long rateToMsConversion;  //1000

	private final AtomicInteger consumedTokens = new AtomicInteger();
	private final AtomicLong lastRefillTime = new AtomicLong(0);

	@Deprecated
	public MargoEurekaRateLimiter() {
		this(TimeUnit.SECONDS);
	}

	public MargoEurekaRateLimiter(TimeUnit averageRateUnit) {
		switch (averageRateUnit) {
		case SECONDS:
			rateToMsConversion = 1000;
			break;
		case MINUTES:
			rateToMsConversion = 60 * 1000;
			break;
		default:
			throw new IllegalArgumentException("TimeUnit of " + averageRateUnit + " is not supported");
		}
	}

	// 提供给外界获取 token 的方法     burstSize 突发速率   averageRate 平均速率
	public boolean acquire(int burstSize, long averageRate) {
		return acquire(burstSize, averageRate, System.currentTimeMillis());
	}

	public boolean acquire(int burstSize, long averageRate, long currentTimeMillis) {
		if (burstSize <= 0 || averageRate <= 0) { // Instead of throwing exception, we just let all the traffic go
			return true;
		}

		refillToken(burstSize, averageRate, currentTimeMillis);
		return consumeToken(burstSize);
	}

	//添加token
	private void refillToken(int burstSize, long averageRate, long currentTimeMillis) {
		long refillTime = lastRefillTime.get();
		long timeDelta = currentTimeMillis - refillTime;
		//根据频率计算需要增加多少 token
		long newTokens = timeDelta * averageRate / rateToMsConversion;
		if (newTokens > 0) {
			long newRefillTime = refillTime == 0 ? currentTimeMillis
					: refillTime + newTokens * rateToMsConversion / averageRate;
			// CAS 保证有且仅有一个线程进入填充
			if (lastRefillTime.compareAndSet(refillTime, newRefillTime)) {
				while (true) {
					int currentLevel = consumedTokens.get();
					int adjustedLevel = Math.min(currentLevel, burstSize); // In case burstSize decreased
					int newLevel = (int) Math.max(0, adjustedLevel - newTokens);
					// while true 直到更新成功为止
					if (consumedTokens.compareAndSet(currentLevel, newLevel)) {
						return;
					}
				}
			}
		}
	}

	//消费token
	private boolean consumeToken(int burstSize) {
		while (true) {
			int currentLevel = consumedTokens.get();
			if (currentLevel >= burstSize) {
				return false;
			}
			if (consumedTokens.compareAndSet(currentLevel, currentLevel + 1)) {
				return true;
			}
		}
	}

	public void reset() {
		consumedTokens.set(0);
		lastRefillTime.set(0);
	}
}
